package com.kh.rationaldrugreview.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.kh.rationaldrugreview.entity.KnowledgeEntity;
import com.kh.rationaldrugreview.entity.EntityRelationship;
import com.kh.rationaldrugreview.neo4j.repository.DrugNodeRepository;
import com.kh.rationaldrugreview.neo4j.service.Neo4jCypherService;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.*;

/**
 * 药品知识图谱构建服务
 * 负责从PDF药品说明书中提取实体和关系，构建知识图谱
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DrugKnowledgeGraphService {
    
    private final ChatClient chatClient;
    private final PDFTextExtractor pdfTextExtractor;
    private final DrugNodeRepository drugNodeRepository;
    private final Neo4jCypherService neo4jCypherService;
    private final ObjectMapper objectMapper;
    private final KnowledgeEntityService knowledgeEntityService;
    
    /**
     * 知识图谱提取结果
     */
    @Data
    public static class KnowledgeGraphResult {
        private List<Entity> entities = new ArrayList<>();
        private List<Relationship> relationships = new ArrayList<>();
        private Map<String, Object> metadata = new HashMap<>();
        private String status;
        private String message;
    }
    
    @Data
    public static class Entity {
        private String id;
        private String type;
        private String name;
        private Map<String, Object> attributes = new HashMap<>();
    }
    
    @Data
    public static class Relationship {
        private String id;
        private String type;
        private String source;
        private String target;
        private Integer strength;
        private Map<String, Object> properties = new HashMap<>();
    }
    
    /**
     * 处理上传的PDF文件并构建知识图谱
     */
    @Transactional
    public KnowledgeGraphResult processPDFToKnowledgeGraph(MultipartFile file) throws IOException {
        log.info("Processing PDF file: {} (size: {} bytes)", 
            file.getOriginalFilename(), file.getSize());
        
        KnowledgeGraphResult result = new KnowledgeGraphResult();
        
        try {
            // 1. 保存临时文件
            File tempFile = saveToTempFile(file);
            
            // 2. 提取PDF文本
            String pdfText = pdfTextExtractor.extractTextFromPDF(tempFile.getAbsolutePath());
            log.info("Extracted {} characters from PDF", pdfText.length());
            
            // 3. 调用LLM提取实体和关系
            String extractedJson = extractEntitiesAndRelationships(pdfText);
            
            // 4. 解析JSON结果
            result = parseExtractionResult(extractedJson);
            
            // 5. 验证和补充提取结果
            result = validateAndEnhance(result, pdfText);
            
            // 6. 保存到Neo4j
            saveToNeo4j(result);
            
            // 7. 推理间接关系
            inferAdditionalRelationships(result);
            
            // 8. 设置元数据
            result.getMetadata().put("filename", file.getOriginalFilename());
            result.getMetadata().put("fileSize", file.getSize());
            result.getMetadata().put("textLength", pdfText.length());
            result.getMetadata().put("processedAt", new Date());
            result.setStatus("SUCCESS");
            result.setMessage("Knowledge graph built successfully");
            
            // 清理临时文件
            tempFile.delete();
            
            log.info("Successfully built knowledge graph with {} entities and {} relationships",
                result.getEntities().size(), result.getRelationships().size());
            
        } catch (Exception e) {
            log.error("Error processing PDF to knowledge graph", e);
            result.setStatus("ERROR");
            result.setMessage("Failed to process PDF: " + e.getMessage());
        }
        
        return result;
    }
    
    /**
     * 保存上传文件到临时目录
     */
    private File saveToTempFile(MultipartFile file) throws IOException {
        String tempDir = System.getProperty("java.io.tmpdir");
        File tempFile = new File(tempDir, "drug_manual_" + 
            System.currentTimeMillis() + "_" + file.getOriginalFilename());
        
        try (FileOutputStream fos = new FileOutputStream(tempFile)) {
            fos.write(file.getBytes());
        }
        
        return tempFile;
    }
    
    /**
     * 使用LLM提取实体和关系
     */
    private String extractEntitiesAndRelationships(String pdfText) {
        log.info("Calling LLM to extract entities and relationships");
        
        // 限制文本长度避免超过token限制
        String truncatedText = pdfText;
        if (pdfText.length() > 20000) {
            truncatedText = pdfText.substring(0, 20000);
            log.warn("Text truncated to 20000 characters to fit token limit");
        }
        
        String prompt = KnowledgeGraphPromptTemplate.ENTITY_EXTRACTION_PROMPT
            .replace("{manualText}", truncatedText);
        
        String response = chatClient.prompt()
            .system(KnowledgeGraphPromptTemplate.SYSTEM_PROMPT)
            .user(prompt)
            .call()
            .content();
        
        log.info("LLM response received, length: {} characters", response.length());
        
        // 提取JSON部分
        return extractJsonFromResponse(response);
    }
    
    /**
     * 从LLM响应中提取JSON
     */
    private String extractJsonFromResponse(String response) {
        // 查找JSON开始和结束位置
        int startIdx = response.indexOf("{");
        int endIdx = response.lastIndexOf("}");
        
        if (startIdx != -1 && endIdx != -1 && endIdx > startIdx) {
            return response.substring(startIdx, endIdx + 1);
        }
        
        // 如果没有找到，尝试查找```json标记
        if (response.contains("```json")) {
            startIdx = response.indexOf("```json") + 7;
            endIdx = response.indexOf("```", startIdx);
            if (endIdx > startIdx) {
                return response.substring(startIdx, endIdx).trim();
            }
        }
        
        return response;
    }
    
    /**
     * 解析提取结果
     */
    private KnowledgeGraphResult parseExtractionResult(String json) {
        try {
            return objectMapper.readValue(json, KnowledgeGraphResult.class);
        } catch (Exception e) {
            log.error("Failed to parse extraction result", e);
            // 返回空结果
            return new KnowledgeGraphResult();
        }
    }
    
    /**
     * 验证和增强提取结果
     */
    private KnowledgeGraphResult validateAndEnhance(KnowledgeGraphResult result, String originalText) {
        log.info("Validating and enhancing extraction result");
        
        try {
            String currentJson = objectMapper.writeValueAsString(result);
            
            String validationPrompt = KnowledgeGraphPromptTemplate.VALIDATION_PROMPT
                .replace("{extractedData}", currentJson)
                .replace("{originalText}", originalText.substring(0, Math.min(5000, originalText.length())));
            
            String enhancedJson = chatClient.prompt()
                .system(KnowledgeGraphPromptTemplate.SYSTEM_PROMPT)
                .user(validationPrompt)
                .call()
                .content();
            
            String cleanJson = extractJsonFromResponse(enhancedJson);
            KnowledgeGraphResult enhancedResult = objectMapper.readValue(cleanJson, KnowledgeGraphResult.class);
            
            log.info("Enhanced result: {} entities, {} relationships",
                enhancedResult.getEntities().size(), enhancedResult.getRelationships().size());
            
            return enhancedResult;
            
        } catch (Exception e) {
            log.error("Failed to validate and enhance", e);
            return result; // 返回原始结果
        }
    }
    
    /**
     * 保存到Neo4j图数据库
     * 使用KnowledgeEntityService进行实体去重
     */
    private void saveToNeo4j(KnowledgeGraphResult result) {
        log.info("Saving {} entities and {} relationships to Neo4j with deduplication",
            result.getEntities().size(), result.getRelationships().size());
        
        // 保存实体，使用MySQL进行去重管理
        // 维护LLM生成的entity ID到数据库canonical entity ID的映射
        Map<String, Long> llmIdToCanonicalId = new HashMap<>();
        Map<String, KnowledgeEntity> canonicalEntities = new HashMap<>();
        
        // 获取source信息（可以从metadata中获取）
        String source = result.getMetadata() != null ? 
            (String) result.getMetadata().get("filename") : "unknown";
        
        for (Entity entity : result.getEntities()) {
            try {
                // 1. 使用KnowledgeEntityService查找或创建规范化的实体
                KnowledgeEntity canonicalEntity = knowledgeEntityService.findOrCreateEntity(
                    entity.getType(),
                    entity.getName(),
                    entity.getAttributes(),
                    source
                );
                
                // 2. 记录LLM ID到canonical ID的映射
                llmIdToCanonicalId.put(entity.getId(), canonicalEntity.getEntityId());
                canonicalEntities.put(entity.getId(), canonicalEntity);
                
                // 3. 准备Neo4j节点属性
                Map<String, Object> properties = new HashMap<>(entity.getAttributes());
                properties.put("name", entity.getName());
                properties.put("canonicalId", canonicalEntity.getEntityId());  // 使用数据库的canonical ID
                properties.put("displayName", canonicalEntity.getDisplayName());
                
                // 处理患者特征相关的特殊实体类型
                if ("PATIENT_PROFILE".equals(entity.getType())) {
                    properties = enrichPatientProfileProperties(properties, entity.getAttributes());
                } else if ("PREGNANCY_CATEGORY".equals(entity.getType())) {
                    properties = enrichPregnancyCategoryProperties(properties, entity.getAttributes());
                } else if ("AGE_GROUP".equals(entity.getType())) {
                    properties = enrichAgeGroupProperties(properties, entity.getAttributes());
                } else if ("ORGAN_FUNCTION".equals(entity.getType())) {
                    properties = enrichOrganFunctionProperties(properties, entity.getAttributes());
                } else if ("METABOLIZER_TYPE".equals(entity.getType())) {
                    properties = enrichMetabolizerTypeProperties(properties, entity.getAttributes());
                }
                
                // 4. 在Neo4j中创建或更新节点
                // 使用canonical ID作为唯一标识符
                Long neo4jNodeId = neo4jCypherService.createOrUpdateEntity(
                    entity.getType(), 
                    String.valueOf(canonicalEntity.getEntityId()),  // 使用canonical ID
                    properties
                );
                
                // 5. 更新MySQL中记录的Neo4j节点ID
                if (neo4jNodeId != null && canonicalEntity.getNeo4jNodeId() == null) {
                    knowledgeEntityService.updateNeo4jNodeId(
                        canonicalEntity.getEntityId(), 
                        String.valueOf(neo4jNodeId)
                    );
                }
                
                log.debug("Processed {} entity: {} (canonical ID: {})", 
                    entity.getType(), entity.getName(), canonicalEntity.getEntityId());
                    
            } catch (Exception e) {
                log.error("Failed to process entity: {}", entity.getName(), e);
            }
        }
        
        // 保存关系，使用canonical实体ID
        for (Relationship rel : result.getRelationships()) {
            try {
                // 1. 获取源和目标的canonical实体
                Long sourceCanonicalId = llmIdToCanonicalId.get(rel.getSource());
                Long targetCanonicalId = llmIdToCanonicalId.get(rel.getTarget());
                
                if (sourceCanonicalId == null || targetCanonicalId == null) {
                    log.warn("Skipping relationship due to missing entities: {} -> {}", 
                        rel.getSource(), rel.getTarget());
                    continue;
                }
                
                // 2. 在MySQL中记录关系（去重）
                EntityRelationship canonicalRel = knowledgeEntityService.findOrCreateRelationship(
                    sourceCanonicalId,
                    targetCanonicalId,
                    rel.getType(),
                    rel.getProperties(),
                    rel.getStrength(),
                    source
                );
                
                // 3. 准备Neo4j关系属性
                Map<String, Object> properties = new HashMap<>(rel.getProperties());
                properties.put("strength", rel.getStrength());
                properties.put("canonicalRelId", canonicalRel.getRelationshipId());
                
                // 处理患者特征相关的特殊关系类型
                if (isPatientRelatedRelationship(rel.getType())) {
                    properties = enrichPatientRelationshipProperties(rel.getType(), properties, rel.getProperties());
                }
                
                // 4. 在Neo4j中创建关系
                // 使用canonical ID作为节点标识
                neo4jCypherService.createRelationship(
                    String.valueOf(sourceCanonicalId),
                    String.valueOf(targetCanonicalId),
                    rel.getType(),
                    properties
                );
                
                log.debug("Created {} relationship: {} -> {} (canonical: {} -> {})",
                    rel.getType(), rel.getSource(), rel.getTarget(),
                    sourceCanonicalId, targetCanonicalId);
                    
            } catch (Exception e) {
                log.error("Failed to create relationship: {}", rel.getId(), e);
            }
        }
        
        log.info("Completed saving to Neo4j with deduplication. Entities processed: {}, Relationships processed: {}",
            llmIdToCanonicalId.size(), result.getRelationships().size());
    }
    
    /**
     * 丰富患者画像属性
     */
    private Map<String, Object> enrichPatientProfileProperties(Map<String, Object> properties, Map<String, Object> attributes) {
        // 确保关键属性存在
        properties.putIfAbsent("ageGroup", attributes.getOrDefault("ageGroup", "成人"));
        properties.putIfAbsent("gender", attributes.getOrDefault("gender", "不限"));
        properties.putIfAbsent("pregnancyStatus", attributes.getOrDefault("pregnancyStatus", "不适用"));
        properties.putIfAbsent("lactationStatus", attributes.getOrDefault("lactationStatus", "不适用"));
        properties.putIfAbsent("renalFunction", attributes.getOrDefault("renalFunction", "正常"));
        properties.putIfAbsent("hepaticFunction", attributes.getOrDefault("hepaticFunction", "正常"));
        
        // 处理年龄范围
        if (attributes.containsKey("ageRangeMin") || attributes.containsKey("ageRangeMax")) {
            properties.put("ageRangeMin", attributes.getOrDefault("ageRangeMin", 0));
            properties.put("ageRangeMax", attributes.getOrDefault("ageRangeMax", 120));
        }
        
        return properties;
    }
    
    /**
     * 丰富妊娠分级属性
     */
    private Map<String, Object> enrichPregnancyCategoryProperties(Map<String, Object> properties, Map<String, Object> attributes) {
        properties.putIfAbsent("category", attributes.getOrDefault("category", "未分级"));
        properties.putIfAbsent("riskLevel", attributes.getOrDefault("riskLevel", "未知"));
        properties.putIfAbsent("description", attributes.getOrDefault("description", ""));
        properties.putIfAbsent("chineseDescription", attributes.getOrDefault("chineseDescription", ""));
        return properties;
    }
    
    /**
     * 丰富年龄组属性
     */
    private Map<String, Object> enrichAgeGroupProperties(Map<String, Object> properties, Map<String, Object> attributes) {
        properties.putIfAbsent("groupName", attributes.getOrDefault("groupName", ""));
        
        // 转换年龄为月份（如果以年为单位）
        if (attributes.containsKey("ageRangeMinYears")) {
            Double minYears = Double.parseDouble(attributes.get("ageRangeMinYears").toString());
            properties.put("ageRangeMin", (int)(minYears * 12));
        }
        if (attributes.containsKey("ageRangeMaxYears")) {
            Double maxYears = Double.parseDouble(attributes.get("ageRangeMaxYears").toString());
            properties.put("ageRangeMax", (int)(maxYears * 12));
        }
        
        return properties;
    }
    
    /**
     * 丰富器官功能属性
     */
    private Map<String, Object> enrichOrganFunctionProperties(Map<String, Object> properties, Map<String, Object> attributes) {
        properties.putIfAbsent("organName", attributes.getOrDefault("organName", ""));
        properties.putIfAbsent("functionLevel", attributes.getOrDefault("functionLevel", "正常"));
        properties.putIfAbsent("clinicalMarkers", attributes.getOrDefault("clinicalMarkers", ""));
        
        // 处理剂量调整系数
        if (attributes.containsKey("dosageAdjustmentFactor")) {
            try {
                Double factor = Double.parseDouble(attributes.get("dosageAdjustmentFactor").toString());
                properties.put("dosageAdjustmentFactor", factor);
            } catch (NumberFormatException e) {
                properties.put("dosageAdjustmentFactor", 1.0);
            }
        }
        
        return properties;
    }
    
    /**
     * 丰富代谢型属性
     */
    private Map<String, Object> enrichMetabolizerTypeProperties(Map<String, Object> properties, Map<String, Object> attributes) {
        properties.putIfAbsent("typeName", attributes.getOrDefault("typeName", "正常代谢型"));
        properties.putIfAbsent("enzyme", attributes.getOrDefault("enzyme", ""));
        properties.putIfAbsent("frequency", attributes.getOrDefault("frequency", ""));
        properties.putIfAbsent("clinicalImplication", attributes.getOrDefault("clinicalImplication", ""));
        properties.putIfAbsent("dosageRecommendation", attributes.getOrDefault("dosageRecommendation", ""));
        return properties;
    }
    
    /**
     * 判断是否为患者相关的关系类型
     */
    private boolean isPatientRelatedRelationship(String relType) {
        return "SUITABLE_FOR".equals(relType) ||
               "CONTRAINDICATED_FOR_PROFILE".equals(relType) ||
               "REQUIRES_ADJUSTMENT_FOR".equals(relType) ||
               "HAS_PREGNANCY_CATEGORY".equals(relType) ||
               "APPROPRIATE_FOR_AGE".equals(relType) ||
               "REQUIRES_ORGAN_FUNCTION".equals(relType) ||
               "AFFECTED_BY_METABOLIZER".equals(relType) ||
               "GENDER_SPECIFIC".equals(relType);
    }
    
    /**
     * 丰富患者相关关系属性
     */
    private Map<String, Object> enrichPatientRelationshipProperties(String relType, Map<String, Object> properties, Map<String, Object> originalProps) {
        switch (relType) {
            case "SUITABLE_FOR":
                properties.putIfAbsent("conditions", originalProps.getOrDefault("conditions", ""));
                properties.putIfAbsent("dosageAdjustment", originalProps.getOrDefault("dosageAdjustment", ""));
                properties.putIfAbsent("monitoringRequired", originalProps.getOrDefault("monitoringRequired", false));
                break;
                
            case "CONTRAINDICATED_FOR_PROFILE":
                properties.putIfAbsent("reason", originalProps.getOrDefault("reason", ""));
                properties.putIfAbsent("severity", originalProps.getOrDefault("severity", "绝对禁忌"));
                properties.putIfAbsent("alternative", originalProps.getOrDefault("alternative", ""));
                break;
                
            case "REQUIRES_ADJUSTMENT_FOR":
                properties.putIfAbsent("adjustmentType", originalProps.getOrDefault("adjustmentType", "剂量"));
                properties.putIfAbsent("adjustmentFactor", originalProps.getOrDefault("adjustmentFactor", 1.0));
                properties.putIfAbsent("monitoringParameters", originalProps.getOrDefault("monitoringParameters", ""));
                break;
                
            case "HAS_PREGNANCY_CATEGORY":
                properties.putIfAbsent("evidenceQuality", originalProps.getOrDefault("evidenceQuality", "中"));
                properties.putIfAbsent("updateDate", originalProps.getOrDefault("updateDate", new Date().toString()));
                properties.putIfAbsent("source", originalProps.getOrDefault("source", "药品说明书"));
                properties.putIfAbsent("trimesterSpecific", originalProps.getOrDefault("trimesterSpecific", ""));
                break;
                
            case "APPROPRIATE_FOR_AGE":
                properties.putIfAbsent("dosageByWeight", originalProps.getOrDefault("dosageByWeight", ""));
                properties.putIfAbsent("dosageByAge", originalProps.getOrDefault("dosageByAge", ""));
                properties.putIfAbsent("maxDose", originalProps.getOrDefault("maxDose", ""));
                properties.putIfAbsent("formulation", originalProps.getOrDefault("formulation", ""));
                break;
                
            case "REQUIRES_ORGAN_FUNCTION":
                properties.putIfAbsent("minimumFunction", originalProps.getOrDefault("minimumFunction", ""));
                properties.putIfAbsent("dosageAdjustment", originalProps.getOrDefault("dosageAdjustment", ""));
                properties.putIfAbsent("monitoringFrequency", originalProps.getOrDefault("monitoringFrequency", ""));
                properties.putIfAbsent("contraindicationThreshold", originalProps.getOrDefault("contraindicationThreshold", ""));
                break;
                
            case "AFFECTED_BY_METABOLIZER":
                properties.putIfAbsent("enzyme", originalProps.getOrDefault("enzyme", ""));
                properties.putIfAbsent("impact", originalProps.getOrDefault("impact", ""));
                properties.putIfAbsent("dosageModification", originalProps.getOrDefault("dosageModification", ""));
                properties.putIfAbsent("testingRecommended", originalProps.getOrDefault("testingRecommended", false));
                break;
                
            case "GENDER_SPECIFIC":
                properties.putIfAbsent("gender", originalProps.getOrDefault("gender", ""));
                properties.putIfAbsent("reason", originalProps.getOrDefault("reason", ""));
                properties.putIfAbsent("dosageDifference", originalProps.getOrDefault("dosageDifference", ""));
                properties.putIfAbsent("efficacyDifference", originalProps.getOrDefault("efficacyDifference", ""));
                break;
        }
        
        return properties;
    }
    
    /**
     * 推理额外的关系
     */
    private void inferAdditionalRelationships(KnowledgeGraphResult result) {
        log.info("Inferring additional relationships");
        
        try {
            String currentJson = objectMapper.writeValueAsString(result);
            
            String inferencePrompt = KnowledgeGraphPromptTemplate.RELATIONSHIP_INFERENCE_PROMPT
                .replace("{knownData}", currentJson);
            
            String inferredJson = chatClient.prompt()
                .system(KnowledgeGraphPromptTemplate.SYSTEM_PROMPT)
                .user(inferencePrompt)
                .call()
                .content();
            
            String cleanJson = extractJsonFromResponse(inferredJson);
            KnowledgeGraphResult inferredRelationships = objectMapper.readValue(
                cleanJson, KnowledgeGraphResult.class);
            
            // 添加推理的关系
            if (inferredRelationships.getRelationships() != null) {
                result.getRelationships().addAll(inferredRelationships.getRelationships());
                log.info("Added {} inferred relationships", 
                    inferredRelationships.getRelationships().size());
            }
            
        } catch (Exception e) {
            log.error("Failed to infer additional relationships", e);
        }
    }
    
    /**
     * 获取知识图谱统计信息
     */
    public Map<String, Object> getGraphStatistics() {
        return neo4jCypherService.getGraphStatistics();
    }
}